library(tidyverse)

# define functions----------

#' Fill non-exist marker gene with expression 0
#'
#' so that following analysis throw no error with NA.
#' 
#' @param meta_df A data frame with columns as genes and rows as cells.
#'
#' @return A data frame with filled 0 expression col, if the genes in `all_marker` not exist in the `meta_df`.
#' @export
#'
#' @examples
#' expr <- fetchData(sobj, c('IGHG1', 'IGHG2'))
#' expr_filled <- fill_expr_mat(expr)
fill_expr_mat <- function(meta_df) {
  if (every(all_marker, \(x)x %in% colnames(meta_df))) meta_df
  
  unexist_marker <- all_marker %>%
    discard(\(x)x %in% colnames(meta_df))
  
  matrix(0, nrow = dim(meta_df)[[1]], ncol = length(unexist_marker)) %>%
    as.data.frame() %>%
    set_names(unexist_marker) %>%
    bind_cols(meta_df) %>%
    as_tibble()
}

annotate_complex <- function(df) {
  df |>
    mutate(
      IGL = IGKC | IGLC1 | IGLC2 | IGLC3 | IGLC6 | IGLC7,
      IGH = IGHA1 | IGHA2 | IGHG1 | IGHG2 | IGHG3 | IGHG4 | IGHD | IGHE | IGHM,
      CD79 = CD79A & CD79B,
      TRB = TRBC1 | TRBC2,
      TCRab = TRAC & TRB,
      TRG = TRGC1 | TRGC2,
      TCRgd = TRG & TRDC,
      TCR = TCRab | TCRgd,
      CD3 = CD3D & CD3E & CD3G & CD247,
      Zxk = IGL & IGH & CD79 & TCR & CD3,
      Ahmed = CD19 & CD5 & TRAC & TRB
    )
}

annotate_complex2 <- function(df) {
  df |>
    mutate(
      IGL = case_when(IGL == 1 ~ 1, .default = IGKC | IGLC2 | IGLC3 | IGLC6 | IGLC7),
      IGH = case_when(IGH == 1 ~ 1, .default = IGHA1 | IGHA2 | IGHG1 | IGHG2 | IGHG3 | IGHG4 | IGHD | IGHE | IGHM),
      CD79 = case_when(CD79 == 1 ~ 1, .default = CD79A & CD79B),
      TRB = case_when(TRB == 1 ~ 1, .default = TRBC1 | TRBC2),
      TCRab = case_when(TCRab == 1 ~ 1, .default = TRAC & TRB),
      TRG = case_when(TRG == 1 ~ 1, .default = TRGC1 | TRGC2),
      TCRgd = case_when(TCRgd == 1 ~ 1, .default = TRG & TRDC),
      TCR = case_when(TCR == 1 ~ 1, .default = TCRab | TCRgd),
      CD3 = case_when(CD3 == 1 ~ 1, .default = CD3D & CD3E & CD3G & CD247),
      Zxk = case_when(Zxk == 1 ~ 1, .default = IGL & IGH & CD79 & TCR & CD3),
      Ahmed = case_when(Ahmed == 1 ~ 1, .default = CD19 & CD5 & TRAC & TRB)
    )
}

count_and_fraction <- function(tbbl, component) {
  tbbl |>
    reframe(
      count = sum({{ component }}),
      fraction = sum({{ component }}) / n(),
      total = n()
    )
}

#' Make an iterable tibble from a character vector
#'
#' @param list A character vector as the colnames of output tibble
#'
#' @return A tibble of colnames the same as `list`
#' @export
#'
#' @examples
make_dummy <- function(list) {
  dummy <- list |>
    set_names(list) |>
    as_tibble_row(.name_repair = "minimal")
}

#' Fill one column with value 1
#'
#' @param tbbl A data frame, row as cells, column as genes
#' @param component Character, name of the column to be filled 
#'
#' @return
#' @export
#'
#' @examples
leave_one_out <- function(tbbl, component) {
  tbbl |>
    mutate("{component}" := 1)
}

leave_two_out <- function(tbbl, var1, var2) {
  tbbl |>
    leave_one_out(var1) |>
    leave_one_out(var2)
}

#' Summarize DE cell fraction (Ahmed definition, Zhang XK definition, and leave-one-outs) in every given group.
#'
#' @param aid A data frame of expression matrix, columns as genes, rows as cells
#' @param group A vector in `aid`, group variable in output summary
#'
#' @return A tibble of DE cell fraction value summary. First column will be the variable `group`. Second and third column is count and fraction of DE cells, respectively. Forth column is total cell count in this group. Fifth column indicate DE cell definition, show a gene name for leave-this-gene-out. 
#' @export
#'
#' @examples
calc_de_cell_frac <- function(aid, group) {
  grouped_aid <- aid |>
    fill_expr_mat() |>
    annotate_complex() |>
    group_by({{ group }})
  
  ahmed_frac <- grouped_aid |>
    count_and_fraction(Ahmed) |>
    add_column(type = "Ahmed DE")
  
  zxk_frac <- grouped_aid |>
    count_and_fraction(Zxk) |>
    add_column(type = "Full DE")
  
  all_frac <- loosing_list |>
    map(\(x)leave_one_out(grouped_aid, x)) |>
    map(annotate_complex2) |>
    map(count_and_fraction, Zxk) |>
    map2(loosing_list, \(x, y)add_column(x, type = y)) |>
    list_rbind() |>
    bind_rows(ahmed_frac, zxk_frac)
  
  all_frac
}

# reading data ----------
all_marker <- c("IGKC","IGLC1","IGLC2","IGLC3","IGLC6","IGLC7","IGHA1","IGHA2","IGHG1","IGHG2","IGHG3","IGHG4", "IGHD","IGHE","IGHM","CD79A","CD79B","TRBC1","TRBC2","TRAC","TRDC","TRGC1","TRGC2","CD3D","CD3E","CD3G","CD247", "CD19","CD5")

loosing_list <- c("CD79A", "CD79B", "CD79", "IGL", "IGH", "CD3D", "CD3E", "CD3G", "CD247", "CD3", "TCR")

loosing_dummy <- loosing_list |>
  make_dummy()

pair_loosing_dummy <- expand_grid(bcr = c("CD79", "IGH", "IGL"),
                                  tcr = c("CD3", "TCR")) |>
  add_case(bcr = c("CD79", "CD79", 'IGH'), tcr = c("IGH", "IGL", "IGL")) |>
  map(make_dummy)
